# -*- coding: utf-8 -*-
"""
===============================
@Author     : Zuo WenTao
@Time       : 2024/6/27 11:21
@Description: 
@Change     : 
@File       : EsClient.py
===============================
"""
from elasticsearch import Elasticsearch, helpers
from typing import List, Dict, Any
from loguru import logger

ES_HOST = ["http://140.143.51.14:9200"]
ES_USERNAME = "elastic"
ES_PASSWORD = "9xGM5YDWmstkTUAf"
ES_INDEX = "unpaywall-20230410*"


class ElasticsearchPool:
    def __init__(self, hosts: List[str], username: str = None, password: str = None):
        self.es = Elasticsearch(
            hosts=hosts,
            http_auth=(username, password) if username and password else None
        )
        logger.success(f"Es -- {ES_HOST} -- 连接成功")

    def create_index(self, index_name: str, settings: Dict[str, Any] = None, mappings: Dict[str, Any] = None):
        """
        创建索引
        :param index_name: 索引名称
        :param settings: 索引设置
        :param mappings: 索引映射
        """
        body = {}
        if settings:
            body['settings'] = settings
        if mappings:
            body['mappings'] = mappings
        return self.es.indices.create(index=index_name, body=body, ignore=400)

    def delete_index(self, index_name: str):
        """
        删除索引
        :param index_name: 索引名称
        """
        return self.es.indices.delete(index=index_name, ignore=[400, 404])

    def index_document(self, index_name: str, document: Dict[str, Any], doc_id: str = None):
        """
        索引文档
        :param index_name: 索引名称
        :param document: 文档数据
        :param doc_id: 文档ID
        """
        return self.es.index(index=index_name, body=document, id=doc_id, refresh=True)

    def bulk_index_documents(self, index_name: str, documents: List[Dict[str, Any]]):
        """
        批量索引文档
        :param index_name: 索引名称
        :param documents: 文档列表
        """
        actions = [
            {
                "_index": index_name,
                "_source": doc,
            }
            for doc in documents
        ]
        helpers.bulk(self.es, actions)

    def get_document(self, index_name: str, doc_id: str):
        """
        获取文档
        :param index_name: 索引名称
        :param doc_id: 文档ID
        """
        return self.es.get(index=index_name, id=doc_id)

    def search_documents(self, index_name: str, query: Dict[str, Any], size: int = 10):
        """
        搜索文档
        :param index_name: 索引名称
        :param query: 查询DSL
        :param size: 返回结果数量
        """
        body = {
            "query": query,
            "size": size
        }
        return self.es.search(index=index_name, body=body)

    def msearch_documents(self, index_name: str, queries: List[Dict[str, Any]], search_type: str = "doi", reflection_type: str = "keyword"):
        """
        批量查询多个DOI
        :param index_name: 索引名称
        :param queries: 查询DSL列表
        :return: 查询结果
        """
        body = []
        for query in queries:
            body.append({'index': index_name})
            body.append({"query": {"term": {search_type: query}}}) if reflection_type == "keyword" else body.append({"query": {"match": {search_type: query}}})

        response = self.es.msearch(body=body)
        # return [
        #     data["hits"]["hits"][0]["_source"] if data["hits"]["hits"] else None
        #     for data in response['responses']
        # ]
        results = []
        for data in response['responses']:
            if data["hits"]["hits"]:
                results.append([hit["_source"] for hit in data["hits"]["hits"]])
            else:
                results.append([])

        return results[0]

    def update_document(self, index_name: str, doc_id: str, update_body: Dict[str, Any]):
        """
        更新文档
        :param index_name: 索引名称
        :param doc_id: 文档ID
        :param update_body: 更新数据
        """
        return self.es.update(index=index_name, id=doc_id, body={"doc": update_body})

    def delete_document(self, index_name: str, doc_id: str):
        """
        删除文档
        :param index_name: 索引名称
        :param doc_id: 文档ID
        """
        return self.es.delete(index=index_name, id=doc_id)

    def count_documents(self, index_name: str, query: Dict[str, Any]):
        """
        统计文档数量
        :param index_name: 索引名称
        :param query: 查询DSL
        """
        return self.es.count(index=index_name, body={"query": query})

    def index_exists(self, index_name: str):
        """
        检查索引是否存在
        :param index_name: 索引名称
        """
        return self.es.indices.exists(index=index_name)

    def refresh_index(self, index_name: str):
        """
        刷新索引
        :param index_name: 索引名称
        """
        return self.es.indices.refresh(index=index_name)

    def close(self):
        """
        关闭连接
        """
        self.es.transport.close()


es_client = ElasticsearchPool(hosts=ES_HOST, username=ES_USERNAME, password=ES_PASSWORD)

if __name__ == '__main__':
    search_str = "Epilepsy after meningitis: A systematic review and meta-analysis"
    #             Rheumatoid meningitis: A systematic review and meta‐analysis
    index_name = ES_INDEX
    print(es_client.msearch_documents(index_name=index_name, queries=[search_str], search_type="title", reflection_type="text"))
    # query = {'match': {'title': search_str}}
    # print(es_client.search_documents(index_name=index_name, query=query, size=10))
